package com.mori.encrypt.spring.boot.core.filter;

import com.mori.encrypt.spring.boot.core.algorithm.EncryptAlgorithm;
import com.mori.encrypt.spring.boot.core.algorithm.impl.AesWithRsaEncryptAlgorithm;
import com.mori.encrypt.spring.boot.core.properties.EncryptConfigProperties;
import com.mori.encrypt.spring.boot.core.wrapper.EncryptionRequestWrapper;
import com.mori.encrypt.spring.boot.core.wrapper.EncryptionResponseWrapper;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.bind.annotation.RequestMethod;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * <p> 加密解密过滤器 </p>
 *
 * @author mori
 * @since 2020/9/29 15:54
 */
public class EncryptionFilter implements Filter {

	private Logger logger = LoggerFactory.getLogger(EncryptionFilter.class);

	/**
	 * 加解密配置信息
	 */
	private EncryptConfigProperties encryptConfig;
	/**
	 * 默认加解密算法
	 */
	private EncryptAlgorithm encryptAlgorithm = new AesWithRsaEncryptAlgorithm();

	public EncryptionFilter(EncryptConfigProperties encryptConfig) {
		this.encryptConfig = encryptConfig;
	}

	@Override
	public void init(FilterConfig filterConfig) throws ServletException {

	}

	@Override
	public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException,
			ServletException {
		HttpServletRequest req = (HttpServletRequest) request;
		String reqMethod = req.getMethod();

		HttpServletResponse resp = (HttpServletResponse) response;
		String uri = req.getRequestURI();
		logger.debug("RequestURI: {}", uri);
		// 未启用
		if (!encryptConfig.getEnable()) {
			chain.doFilter(req, resp);
			return;
		}
		// 是否需要解密
		boolean decryptionStatus = this.contains(encryptConfig.getRequestDecryptUriList(), uri, reqMethod);
		// 是否需要加密
		boolean encryptionStatus = this.contains(encryptConfig.getResponseEncryptUriList(), uri, reqMethod);
		// 没有加解密操作
		if (!decryptionStatus && !encryptionStatus) {
			chain.doFilter(req, resp);
			return;
		}
		EncryptionResponseWrapper responseWrapper = null;
		EncryptionRequestWrapper requestWrapper = null;
		if (decryptionStatus) {
			requestWrapper = new EncryptionRequestWrapper(req);
			// 解密出来
			processDecryption(requestWrapper, req);
		}
		if (encryptionStatus) {
			responseWrapper = new EncryptionResponseWrapper(resp);
		}
		// 同时需要加解密
		if (encryptionStatus && decryptionStatus) {
			chain.doFilter(requestWrapper, responseWrapper);
		} else if (encryptionStatus) {
			// 只需要响应加密
			chain.doFilter(req, responseWrapper);
		} else {
			// 只需要请求解密
			chain.doFilter(requestWrapper, resp);
		}
		if (encryptionStatus) {
			// 响应加密
			String responseData = responseWrapper.getResponseData();
			writeEncryptContent(responseData, response);
		}
	}

	/**
	 * 请求解密处理
	 *
	 * @param requestWrapper
	 * @param req
	 */
	private void processDecryption(EncryptionRequestWrapper requestWrapper, HttpServletRequest req) {
		String requestData = requestWrapper.getRequestData();
		String uri = req.getRequestURI();
		logger.debug("RequestData: {}", requestData);
		try {
			if (!StringUtils.endsWithIgnoreCase(req.getMethod(), RequestMethod.GET.name())) {
				// 解密请求体
				String decryptRequestData = encryptAlgorithm.decryptRequestBody(requestData, encryptConfig.getBodyDecryptKey());
				logger.debug("DecryptRequestData: {}", decryptRequestData);
				requestWrapper.setRequestData(decryptRequestData);
			}
			// url参数解密
			Map<String, String> paramMap = new HashMap<>();
			Enumeration<String> parameterNames = req.getParameterNames();
			while (parameterNames.hasMoreElements()) {
				String paramName = parameterNames.nextElement();
				String prefixUri = req.getMethod().toLowerCase() + ":" + uri;
				if (encryptConfig.getRequestDecryptParams(prefixUri).contains(paramName)) {
					String paramValue = req.getParameter(paramName);
					String decryptParamValue = encryptAlgorithm.decrypt(paramValue, encryptConfig.getParamDecryptKey());
					paramMap.put(paramName, decryptParamValue);
				}
			}
			requestWrapper.setParamMap(paramMap);
		} catch (Exception e) {
			logger.error("请求数据解密失败", e);
			throw new RuntimeException(e);
		}
	}

	/**
	 * 输出加密内容
	 *
	 * param responseData
	 * param response
	 * @throws IOException*/
	private void writeEncryptContent(String responseData, ServletResponse response) throws IOException {
		logger.debug("ResponseData: {}", responseData);
		ServletOutputStream out = null;
		try {
			responseData = encryptAlgorithm.encrypt(responseData, encryptConfig.getEncryptKey());
			logger.debug("EncryptResponseData: {}", responseData);
			response.setContentLength(responseData.length());
			response.setCharacterEncoding(encryptConfig.getResponseCharset());
			out = response.getOutputStream();
			out.write(responseData.getBytes(encryptConfig.getResponseCharset()));
		} catch (Exception e) {
			logger.error("响应数据加密失败", e);
			throw new RuntimeException(e);
		} finally {
			if (out != null) {
				out.flush();
				out.close();
			}
		}
	}

	private boolean contains(List<String> list, String uri, String methodType) {
		if (list.contains(uri)) {
			return true;
		}
		String prefixUri = methodType.toLowerCase() + ":" + uri;
		logger.debug("contains uri: {}", prefixUri);
		return list.contains(prefixUri);
	}

	@Override
	public void destroy() {

	}

}
